import asyncio
import aiohttp
import json

from django.conf import settings

from framework.exceptions import TDEngineRequestException


class TDEngine:
    """
    异步 请求 http request tdengine
    """

    def __init__(self, conf, f=None):
        self.f = f
        self.conf = conf

    @staticmethod
    def fmt_response(response: dict, request_json: dict):
        """对返回格式进行格式化处理"""
        if response.get('code') != 0:
            raise TDEngineRequestException('request_json > %s.' % json.dumps(request_json))
        return response['data']

    async def request(self, session: aiohttp.ClientSession, sql: str):
        """
        3x return {'code': 4357, 'desc': 'no auth info input', 'timing': 96458}
        :param session:
        :param sql:
        :return:
        """
        async with session.post(self.conf['url'], data=sql) as response:
            if self.f:
                return self.f(await response.json(), sql)
            return await response.json()

    async def main(self, sqls: list):
        async with aiohttp.ClientSession(
                auth=aiohttp.BasicAuth(self.conf['username'], self.conf['password']),
                timeout=aiohttp.ClientTimeout(self.conf.get('timeout', 10))
        ) as session:
            tasks = [self.request(session, item) for item in sqls]
            return await asyncio.gather(*tasks)

    def run(self, *args, **kwargs):
        """函数入口"""
        return asyncio.run(self.main(*args, **kwargs))


td_engine = TDEngine(conf=settings.TDENGINE)

if __name__ == '__main__':
    ret = td_engine.main(['select ts, url, time from log_2024;'])
    print(ret)
    ...
